{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bfd62b9362ec4dbb825e786821358a5e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(layout=Layout(height='1in'))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**tacotron2 input**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fead313d50594d0688a399cff8b6eb86",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Textarea(value='type here', layout=Layout(height='80px', width='550px'), placeholder='')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**tacotron2 preprocessing**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e610d8bd77ad44a7839d6ede103a71ef",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output(layout=Layout(height='1in', object_fit='fill', object_position='{center} {center}', width='10in'))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**tacotron2 output / waveglow input**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "240214dbb35c4bfd97b84df7cb5cd8cf",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output(layout=Layout(height='2.1in', object_fit='fill', object_position='{center} {center}', width='10in'))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**waveglow output**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5bc0e4519ea74a89b34c55fa8a040b02",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output(layout=Layout(height='2in', object_fit='fill', object_position='{center} {center}', width='10in'))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**play**"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "30976ec770c34d34880af8d7f3d6ff08",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output(layout=Layout(height='1in', object_fit='fill', object_position='{center} {center}', width='10in'))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bfd62b9362ec4dbb825e786821358a5e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(layout=Layout(height='1in'))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import os\n",
    "import time\n",
    "import numpy as np\n",
    "import collections\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import cm as cm\n",
    "from IPython.display import Audio, display, clear_output, Markdown, Image\n",
    "import librosa\n",
    "import librosa.display\n",
    "import ipywidgets as widgets\n",
    "# \n",
    "from tacotron2.text import text_to_sequence as text_to_sequence_internal\n",
    "from tacotron2.text.symbols import symbols\n",
    "# \n",
    "from tensorrtserver.api import *\n",
    "\n",
    "\n",
    "defaults = {\n",
    "    # settings\n",
    "    'sigma_infer': 0.6,        # don't touch this\n",
    "    'sampling_rate': 22050,    # don't touch this\n",
    "    'stft_hop_length': 256,    # don't touch this\n",
    "    'url': 'localhost:8000',   # don't touch this\n",
    "    'protocol': 0,             # 0: http, 1: grpc \n",
    "    'autoplay': True,          # autoplay\n",
    "    'character_limit_min': 4,  # don't touch this\n",
    "    'character_limit_max': 124 # don't touch this\n",
    "}\n",
    "\n",
    "\n",
    "# create args object\n",
    "class Struct:\n",
    "    def __init__(self, **entries):\n",
    "        self.__dict__.update(entries)\n",
    "\n",
    "\n",
    "args = Struct(**defaults)\n",
    "\n",
    "\n",
    "# create the inference context for the models\n",
    "infer_ctx_tacotron2 = InferContext(args.url, args.protocol, 'tacotron2', -1)\n",
    "infer_ctx_waveglow = InferContext(args.url, args.protocol, 'waveglow', -1)\n",
    "\n",
    "\n",
    "def display_heatmap(sequence, title='preprocessed text'):\n",
    "    ''' displays sequence as a heatmap '''\n",
    "    clear_output(wait=True)\n",
    "    sequence = sequence[None, :]\n",
    "    plt.figure(figsize=(10, 2.5))\n",
    "    plt.title(title)\n",
    "    plt.tick_params(\n",
    "        axis='both',\n",
    "        which='both',\n",
    "        bottom=False,\n",
    "        top=False,\n",
    "        left=False,\n",
    "        right=False,\n",
    "        labelbottom=False,\n",
    "        labelleft=False)\n",
    "    plt.imshow(sequence, cmap='BrBG_r', interpolation='nearest')\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def display_sound(signal, title, color):\n",
    "    ''' displays signal '''\n",
    "    clear_output(wait=True)\n",
    "    plt.figure(figsize=(10, 2.5))\n",
    "    plt.title(title)\n",
    "    plt.tick_params(\n",
    "        axis='both',\n",
    "        which='both',\n",
    "        bottom=True,\n",
    "        top=False,\n",
    "        left=False,\n",
    "        right=False,\n",
    "        labelbottom=True,\n",
    "        labelleft=False)\n",
    "    librosa.display.waveplot(signal, color=color)\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def display_spectrogram(mel, title):\n",
    "    ''' displays mel spectrogram '''\n",
    "    clear_output(wait=True)\n",
    "    fig = plt.figure(figsize=(10, 2.5))\n",
    "    ax = fig.add_subplot(111)\n",
    "    plt.title(title)\n",
    "    plt.tick_params(\n",
    "        axis='both',\n",
    "        which='both',\n",
    "        bottom=True,\n",
    "        top=False,\n",
    "        left=False,\n",
    "        right=False,\n",
    "        labelbottom=True,\n",
    "        labelleft=False)\n",
    "    plt.xlabel('Time')\n",
    "    cmap = cm.get_cmap('jet', 30)\n",
    "    cax = ax.imshow(mel.astype(np.float32), interpolation=\"nearest\", cmap=cmap)\n",
    "    ax.grid(True)\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def text_to_sequence(text):\n",
    "    ''' preprocessor of tacotron2\n",
    "        ::text:: the input str\n",
    "        ::returns:: sequence, the preprocessed text\n",
    "    '''\n",
    "    sequence = text_to_sequence_internal(text, ['english_cleaners'])\n",
    "    sequence = np.array(sequence, dtype=np.int64)\n",
    "    return sequence\n",
    "\n",
    "\n",
    "def sequence_to_mel(sequence):\n",
    "    ''' calls tacotron2\n",
    "        ::sequence:: int64 numpy array, contains the preprocessed text\n",
    "        ::returns:: (mel, mel_lengths) pair\n",
    "                     mel is the mel-spectrogram, np.array\n",
    "                     mel_lengths contains the length of the unpadded mel, np.array\n",
    "    '''\n",
    "    input_lengths = [len(sequence)]\n",
    "    input_lengths = np.array(input_lengths, dtype=np.int64)\n",
    "    # prepare input/output\n",
    "    input_dict = {}\n",
    "    input_dict['sequence__0'] = (sequence,)\n",
    "    input_dict['input_lengths__1'] = (input_lengths,)\n",
    "    output_dict = {}\n",
    "    output_dict['mel_outputs_postnet__0'] = InferContext.ResultFormat.RAW\n",
    "    output_dict['mel_lengths__1'] = InferContext.ResultFormat.RAW\n",
    "    batch_size = 1\n",
    "    # call tacotron2\n",
    "    result = infer_ctx_tacotron2.run(input_dict, output_dict, batch_size)\n",
    "    # get results\n",
    "    mel = result['mel_outputs_postnet__0'][0] # take only the first instance in the output batch\n",
    "    mel_lengths = result['mel_lengths__1'][0] # take only the first instance in the output batch\n",
    "    return mel, mel_lengths\n",
    "\n",
    "\n",
    "def force_to_shape(mel, length):\n",
    "    ''' preprocessor of waveglow\n",
    "        :: mel :: numpy array \n",
    "        :: length :: int \n",
    "        :: return :: m padded (or trimmed) to length in dimension 1\n",
    "    '''\n",
    "    diff = length - mel.shape[1]\n",
    "    if 0 < diff:\n",
    "        # pad it\n",
    "        min_value = mel.min()\n",
    "        shape = ((0,0),(0,diff))\n",
    "        ret = np.pad(mel, shape, mode='constant', constant_values=min_value)\n",
    "    else:\n",
    "        # trim it\n",
    "        ret = mel[:,:length]\n",
    "    ret = ret[:,:,None]\n",
    "    return ret\n",
    "\n",
    "\n",
    "def mel_to_signal(mel, mel_lengths):\n",
    "    ''' calls waveglow\n",
    "        ::mel:: mel spectrogram\n",
    "        ::mel_lengths:: original length of mel spectrogram\n",
    "        ::returns:: waveform\n",
    "    '''\n",
    "    # padding/trimming mel to dimension 620\n",
    "    mel = force_to_shape(mel, 620)\n",
    "    # prepare input/output\n",
    "    mel = mel[None,:,:]\n",
    "    input_dict = {}\n",
    "    input_dict['0'] = (mel,)\n",
    "    shape = (8,19840,1)\n",
    "    shape = (1,*shape)\n",
    "    input_dict['1'] = np.random.normal(0.0, 1.0, shape).astype(mel.dtype)\n",
    "    input_dict['1'] = (input_dict['1'],)\n",
    "    output_dict = {}\n",
    "    output_dict['1991'] = InferContext.ResultFormat.RAW\n",
    "    batch_size = 1\n",
    "    # call waveglow\n",
    "    result = infer_ctx_waveglow.run(input_dict, output_dict, batch_size)\n",
    "    # get the results\n",
    "    signal = result['1991'][0] # take only the first instance in the output batch\n",
    "    signal = signal[0] # remove this line, when waveglow supports dynamic batch sizes\n",
    "    # postprocessing of waveglow: trimming signal to its actual size\n",
    "    trimmed_length = mel_lengths[0] * args.stft_hop_length\n",
    "    signal = signal[:trimmed_length] # trim\n",
    "    signal = signal.astype(np.float32)\n",
    "    return signal\n",
    "\n",
    "\n",
    "# widgets\n",
    "def get_output_widget(width, height):\n",
    "    ''' creates an output widget with default values and returns it '''\n",
    "    layout = widgets.Layout(width=width,\n",
    "                            height=height,\n",
    "                            object_fit='fill',\n",
    "                            object_position = '{center} {center}')\n",
    "    ret = widgets.Output(layout=layout)\n",
    "    return ret\n",
    "\n",
    "\n",
    "text_area = widgets.Textarea(\n",
    "    value='type here',\n",
    "    placeholder='',\n",
    "    description='',\n",
    "    disabled=False,\n",
    "    continuous_update=True,\n",
    "    layout=widgets.Layout(width='550px', height='80px')\n",
    ")\n",
    "\n",
    "\n",
    "plot_text_area_preprocessed = get_output_widget(width='10in',height='1in')\n",
    "plot_spectrogram = get_output_widget(width='10in',height='2.1in')\n",
    "plot_signal = get_output_widget(width='10in',height='2.1in')\n",
    "plot_play = get_output_widget(width='10in',height='1in')\n",
    "\n",
    "\n",
    "def text_area_change(change):\n",
    "    ''' this gets called each time text_area.value changes '''\n",
    "    text = change['new']\n",
    "    text = text.strip(' ')\n",
    "    length = len(text)\n",
    "    if length < args.character_limit_min: # too short text\n",
    "        return\n",
    "    if length > args.character_limit_max: # too long text\n",
    "        text_area.value = text[:args.character_limit_max]\n",
    "        return\n",
    "    # preprocess tacotron2\n",
    "    sequence = text_to_sequence(text)\n",
    "    with plot_text_area_preprocessed:\n",
    "        display_heatmap(sequence)\n",
    "    # run tacotron2\n",
    "    mel, mel_lengths = sequence_to_mel(sequence)\n",
    "    with plot_spectrogram:\n",
    "        display_spectrogram(mel, change['new'])\n",
    "    # run waveglow\n",
    "    signal = mel_to_signal(mel, mel_lengths)\n",
    "    with plot_signal:\n",
    "        display_sound(signal, change['new'], 'green')\n",
    "    with plot_play:\n",
    "        clear_output(wait=True)\n",
    "        display(Audio(signal, rate=args.sampling_rate, autoplay=args.autoplay))\n",
    "        # related issue: https://github.com/ipython/ipython/issues/11316\n",
    "\n",
    "\n",
    "# setup callback\n",
    "text_area.observe(text_area_change, names='value')\n",
    "\n",
    "# decorative widgets\n",
    "empty = widgets.VBox([], layout=widgets.Layout(height='1in'))\n",
    "markdown_4 = Markdown('**tacotron2 input**')\n",
    "markdown_5 = Markdown('**tacotron2 preprocessing**')\n",
    "markdown_6 = Markdown('**tacotron2 output / waveglow input**')\n",
    "markdown_7 = Markdown('**waveglow output**')\n",
    "markdown_8 = Markdown('**play**')\n",
    "\n",
    "# display widgets\n",
    "display(\n",
    "    empty, \n",
    "    markdown_4, text_area, \n",
    "#    markdown_5, plot_text_area_preprocessed, \n",
    "    markdown_6, plot_spectrogram, \n",
    "    markdown_7, plot_signal, \n",
    "    markdown_8, plot_play, \n",
    "    empty\n",
    ")\n",
    "\n",
    "# default text\n",
    "text_area.value = \"I think grown-ups just act like they know what they're doing. \""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
